# db_utils.py
import pymysql
import logging
from contextlib import contextmanager

logger = logging.getLogger(__name__)

class DBUtils:
    def __init__(self, host, port, user, password, database):
        self.db_config = {
            'host': host,
            'port': port,
            'user': user,
            'password': password,
            'database': database,
            'charset': 'utf8mb4',
            'cursorclass': pymysql.cursors.DictCursor
        }
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接"""
        conn = pymysql.connect(**self.db_config)
        try:
            yield conn
        finally:
            conn.close()
    
    def execute_query(self, sql, params=None):
        """执行查询并返回结果"""
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                return cursor.fetchall()
    
    def execute_update(self, sql, params=None):
        """执行更新操作"""
        with self.get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                conn.commit()
                return cursor.rowcount
    
    def check_user_exists(self, username):
        """检查用户是否存在"""
        sql = "SELECT * FROM users WHERE username = %s"
        result = self.execute_query(sql, (username,))
        return len(result) > 0
    
    def delete_user(self, username):
        """删除指定用户"""
        sql = "DELETE FROM users WHERE username = %s"
        return self.execute_update(sql, (username,))
    
    def create_test_user(self, username, password, name="测试用户"):
        """创建测试用户"""
        sql = """
        INSERT INTO users (username, password, name, role) 
        VALUES (%s, %s, %s, 'student')
        """
        return self.execute_update(sql, (username, password, name))